#include <stdlib.h>
#include <stdio.h>
#include "AVL.h"

void update_height(AVLNode *root)
{
    int maxh;
    // no child
    if (root->left == NULL && root->right == NULL)
    {
        maxh = -1;
        root->height = maxh + 1;
        return;
    }
    // only left child
    if (root->right == NULL)
    {
        maxh = root->left->height;
        root->height = maxh + 1;
        return;
    }
    // only right child
    if (root->left == NULL)
    {
        maxh = root->right->height;
        root->height = maxh + 1;
        return;
    }
    // two children
    if (root->left->height >= root->right->height)
        maxh = root->left->height;
    else
        maxh = root->right->height;
    root->height = maxh + 1;
    return;
}

void rightRotate(AVLNode **proot)
{
    AVLNode *k2 = *proot;
    AVLNode *k1 = k2->left;
    AVLNode *Y = k1->right;
    *proot = k1;
    k1->right = k2;
    k2->left = Y;
    update_height(k2);
    update_height(k1);
    return;
}

void leftRotate(AVLNode **proot)
{
    AVLNode *k2 = *proot;
    AVLNode *k1 = k2->right;
    AVLNode *Y = k1->left;
    *proot = k1;
    k1->left = k2;
    k2->right = Y;
    update_height(k2);
    update_height(k1);
    return;
}

void rebalance(AVLNode **proot)
{
    AVLNode *root = *proot;

    int flag_l = 0;
    int flag_r = 0;
    int left_height, right_height;
    if (root->left == NULL)
    {
        left_height = -1;
        flag_l = 1;
    }
    else
        left_height = root->left->height;
    if (root->right == NULL)
    {
        right_height = -1;
        flag_r = 1;
    }
    else
        right_height = root->right->height;
    // if balanced
    if (abs(left_height - right_height) <= 1)
        return;
    // has two children
    if (flag_l == 0 && flag_r == 0)
    {
        // left left case
        if (root->left->left->height == root->right->height + 1)
        {
            rightRotate(proot);
            return;
        }
        // left right case
        if (root->left->right->height == root->right->height + 1)
        {
            leftRotate(&(root->left));
            rightRotate(proot);
            return;
        }
        // right right case
        if (root->right->right->height == root->left->height + 1)
        {
            leftRotate(proot);
            return;
        }
        // right left case
        if (root->right->left->height == root->left->height + 1)
        {
            rightRotate(&(root->right));
            leftRotate(proot);
            return;
        }
    }
    // only left child
    else if (flag_l == 0 && flag_r == 1)
    {
        // left left case
        if (root->left->left != NULL)
        {
            rightRotate(proot);
            return;
        }
        // left right case
        if (root->left->right != NULL)
        {
            leftRotate(&(root->left));
            rightRotate(proot);
            return;
        }
    }
    // only right child
    else if (flag_r == 0 && flag_l == 1)
    {
        // right right case
        if (root->right->right != NULL)
        {
            leftRotate(proot);
            return;
        }
        // right left case
        if (root->right->left != NULL)
        {
            rightRotate(&(root->right));
            leftRotate(proot);
            return;
        }
    }
}

AVLNode *insertNode(AVLNode **proot, int x)
{
    AVLNode *root = *proot;

    if (root == NULL)
    {
        // create new node
        AVLNode *newRoot = (AVLNode *)malloc(sizeof(AVLNode));
        newRoot->key = x;
        newRoot->height = 0;
        newRoot->left = NULL;
        newRoot->right = NULL;
        return *proot = newRoot;
    }

    // duplicated key, do nothing
    if (root->key == x)
        return NULL;

    AVLNode *newNode;
    if (root->key > x)
        newNode = insertNode(&(root->left), x); // x is smaller than root, insert left
    else
        newNode = insertNode(&(root->right), x); // x is bigger than root, insert right

    update_height(root);
    rebalance(proot);
    return newNode;
}

AVLNode *deleteMin(AVLNode **proot)
{
    AVLNode *root = *proot;
    AVLNode *minNode;

    // no left child, it's the min node
    if (root->left == NULL)
    {
        minNode = root;
        *proot = root->right;
        return minNode;
    }
    minNode = deleteMin(&(root->left));
    update_height(root);
    rebalance(proot);
    return minNode;
}

AVLNode *deleteRoot(AVLNode **proot)
{
    AVLNode *root = *proot;
    // remove and return root
    AVLNode *currNode = *proot;
    if (root->left == NULL)
    {
        *proot = root->right;
        return currNode;
    }
    if (root->right == NULL)
    {
        *proot = root->left;
        return currNode;
    }
    // root has two children
    AVLNode *minNode = deleteMin(&(root->right));
    minNode->left = root->left;
    minNode->right = root->right;
    *proot = minNode;
    return currNode;
}

AVLNode *deleteNode(AVLNode **proot, int x)
{
    AVLNode *root = *proot;
    AVLNode *matchNode;
    // no such node
    if (root == NULL)
        return NULL;
    if (root->key > x)
        matchNode = deleteNode(&(root->left), x);
    else if (root->key < x)
        matchNode = deleteNode(&(root->right), x);
    else
        matchNode = deleteRoot(proot);
    update_height(*proot);
    rebalance(proot);
    return matchNode;
}

AVLNode *findNode(AVLNode *root, int x)
{
    // no such node
    if (root == NULL)
        return NULL;

    // base case, node found
    if (root->key == x)
        return root;

    // search subtree
    if (root->key > x)
        return findNode(root->left, x);
    return findNode(root->right, x);
}

void destroyTree(AVLNode **proot)
{
    AVLNode *root = *proot;

    // no such node
    if (root == NULL)
        return;

    // if node has child
    if (root->left != NULL || root->right != NULL)
    {
        destroyTree(&(root->left));
        destroyTree(&(root->right));
    }

    // if node is a leaf
    free(root);
    *proot = NULL;
    return;
}